from domain.model.platform.base import BasePlatformChatModel


class GoogleChatModel(BasePlatformChatModel):
    PLATFORM_NAME: str = "google"
    SUPPORT_MODEL_LIST: list[str] = ['gemini-2.5-pro-preview-05-06', 'gemini-2.5-flash-preview-05-20']
    MODEL_INSTANCE_MAP = {}

    def init_model(self, *model_names: str, **options: object) -> "BasePlatformChatModel":
        super().init_model(*model_names, **options)
        from langchain_google_genai import ChatGoogleGenerativeAI
        for model in model_names:
            self.MODEL_INSTANCE_MAP[model] = ChatGoogleGenerativeAI(
                model=model,
                google_api_key=options.get('api_key'),
                temperature=options.get('temperature', 0),
                max_output_tokens=options.get('max_token', 8192),
            )
        return self


google_chat_model = GoogleChatModel()
